{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pKIoAhqZXkbt"
      },
      "source": [
        "# Learn OpenAI Whisper - Chapter 6\n",
        "## Notebook 3: Video Subtitle Generation using Whisper and OpenVINO™\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1e8ZqjqmY7ue2ynLVTDTJXxUCoTQ9XVJ5)\n",
        "\n",
        "In this advanced tutorial, we will leverage the power of OpenAI's Whisper model in conjunction with OpenVINO toolkit to automatically generate subtitles for a sample video. The process will be broken down into the following key steps:\n",
        "\n",
        "1. Obtaining the pre-trained Whisper model\n",
        "2. Setting up the PyTorch model pipeline\n",
        "3. Transforming the model into OpenVINO Intermediate Representation (IR) format using the model conversion API\n",
        "4. Executing the Whisper pipeline with the converted OpenVINO models to generate the subtitles\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z3Htz6uLXkby"
      },
      "source": [
        "## Setting Up the Environment\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1rSRffHmXkby"
      },
      "source": [
        "We start by importing a helper Python utility module called utils.py from our GitHub repository."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget -nv \"https://github.com/PacktPublishing/Learn-OpenAI-Whisper/raw/main/Chapter06/utils.py\" -O utils.py"
      ],
      "metadata": {
        "id": "5IxkRGGXZHwG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we install critical software dependencies to enable working with AI models and speech data."
      ],
      "metadata": {
        "id": "X8NCIh7jffKd"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": [],
        "id": "R9AJOshJXkby"
      },
      "outputs": [],
      "source": [
        "%pip install -q cohere openai tiktoken\n",
        "%pip install -q \"openvino>=2023.1.0\"\n",
        "%pip install -q \"python-ffmpeg<=1.0.16\" moviepy transformers --extra-index-url https://download.pytorch.org/whl/cpu\n",
        "%pip install -q \"git+https://github.com/garywu007/pytube.git\"\n",
        "%pip install -q gradio\n",
        "%pip install -q \"openai-whisper==20231117\" --extra-index-url https://download.pytorch.org/whl/cpu"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tm1KBIgAXkb0"
      },
      "source": [
        "## Initializing the Whisper Model\n",
        "\n",
        "OpenAI's Whisper is a powerful Transformer-based encoder-decoder model, also known as a sequence-to-sequence model, designed for speech recognition tasks. It operates by mapping a sequence of audio spectrogram features to a corresponding sequence of text tokens. The process can be broken down into three main steps:\n",
        "\n",
        "1. **Feature Extraction**: The raw audio inputs are first converted into a log-Mel spectrogram representation using a feature extractor module.\n",
        "\n",
        "2. **Encoding**: The Transformer encoder then processes the spectrogram, generating a sequence of hidden states that capture the essential information from the audio input.\n",
        "\n",
        "3. **Decoding**: Finally, the decoder autoregressively predicts the text tokens, conditioned on both the previously generated tokens and the encoder's hidden states.\n",
        "\n",
        "The architecture of the Whisper model is illustrated in the diagram below:\n",
        "\n",
        "![whisper_architecture.svg](https://user-images.githubusercontent.com/29454499/204536571-8f6d8d77-5fbd-4c6d-8e29-14e734837860.svg)\n",
        "\n",
        "*Source: https://openai.com/research/whisper*\n",
        "\n",
        "By leveraging this powerful architecture, Whisper achieves state-of-the-art performance on various speech recognition benchmarks, making it an ideal choice for our subtitle generation task.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BXNNEelqXkb0"
      },
      "source": [
        "The creators of Whisper have trained several models with varying sizes and capabilities to cater to different use cases and resource constraints. For the purpose of this tutorial, we will be using the `base` model, which offers a good balance between performance and efficiency. However, it's important to note that the steps and techniques demonstrated in this notebook can be easily applied to other models within the Whisper family, allowing you to experiment with different configurations and find the one that best suits your specific requirements."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pX-7rKhCXkb0"
      },
      "outputs": [],
      "source": [
        "from whisper import _MODELS\n",
        "import ipywidgets as widgets\n",
        "\n",
        "model_id = widgets.Dropdown(\n",
        "    options=list(_MODELS),\n",
        "    value='base',\n",
        "    description='Model:',\n",
        "    disabled=False,\n",
        ")\n",
        "\n",
        "model_id"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HTBnlzWzXkb1"
      },
      "outputs": [],
      "source": [
        "import whisper\n",
        "\n",
        "# model = whisper.load_model(model_id.value)\n",
        "model = whisper.load_model(model_id.value, \"cpu\")\n",
        "model.eval()\n",
        "pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ivt20ZOuXkb2"
      },
      "source": [
        "### Converting the Model to OpenVINO Intermediate Representation (IR) Format\n",
        "\n",
        "To achieve optimal performance and efficiency with the OpenVINO toolkit, it is highly recommended to convert the Whisper model into the OpenVINO-specific Intermediate Representation (IR) format. This process requires two key components:\n",
        "\n",
        "1. An initialized model object\n",
        "2. Sample input data for shape inference\n",
        "\n",
        "We will leverage the `ov.convert_model` function provided by OpenVINO to perform the model conversion. This function takes the initialized model object and sample inputs as arguments and returns an OpenVINO-compatible model that is ready to be loaded onto the target device for inference.\n",
        "\n",
        "Once the conversion is complete, we can save the OpenVINO model to disk using the `ov.save_model` function. This allows us to reuse the converted model in future sessions without the need to repeat the conversion process, saving valuable time and resources.\n",
        "\n",
        "By converting the Whisper model to OpenVINO IR format, we can take full advantage of the performance optimizations and hardware acceleration capabilities offered by the OpenVINO toolkit, ensuring efficient and high-quality subtitle generation.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sHqv0rcDXkb2"
      },
      "source": [
        "### Converting the Whisper Encoder to OpenVINO IR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false
        },
        "id": "iqoRzYnAXkb2"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "\n",
        "WHISPER_ENCODER_OV = Path(f\"whisper_{model_id.value}_encoder.xml\")\n",
        "WHISPER_DECODER_OV = Path(f\"whisper_{model_id.value}_decoder.xml\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "An example input is created using a tensor of zeros. The ov.convert_model function is then used to convert the encoder model to OpenVINO's IR format. The converted model is saved to disk for future use."
      ],
      "metadata": {
        "id": "wuGuzFCjsEyT"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hbtJtI0CXkb2"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import openvino as ov\n",
        "\n",
        "mel = torch.zeros((1, 80 if 'v3' not in model_id.value else 128, 3000))\n",
        "audio_features = model.encoder(mel)\n",
        "if not WHISPER_ENCODER_OV.exists():\n",
        "    encoder_model = ov.convert_model(model.encoder, example_input=mel)\n",
        "    ov.save_model(encoder_model, WHISPER_ENCODER_OV)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xug44Bx7Xkb3"
      },
      "source": [
        " ### Converting the Whisper Decoder to OpenVINO IR\n",
        "\n",
        "The Whisper decoder employs a technique called attention caching to reduce computational complexity and improve efficiency. This involves storing the key and value projections from previous steps in the attention modules, which can then be reused in subsequent computations. However, to ensure accurate tracing and conversion of the decoder to OpenVINO IR format, we need to modify this caching mechanism.\n",
        "\n",
        "In the following code cells, we will define custom forward functions for the decoder's attention modules and residual blocks. These modified functions will explicitly handle the caching and retrieval of key and value projections, making the caching process more transparent and traceable.\n",
        "\n",
        "By adapting the decoder's architecture to be more compatible with the OpenVINO conversion process, we can successfully convert the Whisper decoder to OpenVINO IR format, enabling us to leverage the performance benefits of the OpenVINO toolkit while maintaining the decoder's functionality and efficiency."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yHZ-eQX4Xkb3"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from typing import Optional, Tuple\n",
        "from functools import partial\n",
        "\n",
        "\n",
        "def attention_forward(\n",
        "        attention_module,\n",
        "        x: torch.Tensor,\n",
        "        xa: Optional[torch.Tensor] = None,\n",
        "        mask: Optional[torch.Tensor] = None,\n",
        "        kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,\n",
        "):\n",
        "    \"\"\"\n",
        "    Override for forward method of decoder attention module with storing cache values explicitly.\n",
        "    Parameters:\n",
        "      attention_module: current attention module\n",
        "      x: input token ids.\n",
        "      xa: input audio features (Optional).\n",
        "      mask: mask for applying attention (Optional).\n",
        "      kv_cache: dictionary with cached key values for attention modules.\n",
        "      idx: idx for search in kv_cache.\n",
        "    Returns:\n",
        "      attention module output tensor\n",
        "      updated kv_cache\n",
        "    \"\"\"\n",
        "    q = attention_module.query(x)\n",
        "\n",
        "    if xa is None:\n",
        "        # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;\n",
        "        # otherwise, perform key/value projections for self- or cross-attention as usual.\n",
        "        k = attention_module.key(x)\n",
        "        v = attention_module.value(x)\n",
        "        if kv_cache is not None:\n",
        "            k = torch.cat((kv_cache[0], k), dim=1)\n",
        "            v = torch.cat((kv_cache[1], v), dim=1)\n",
        "        kv_cache_new = (k, v)\n",
        "    else:\n",
        "        # for cross-attention, calculate keys and values once and reuse in subsequent calls.\n",
        "        k = attention_module.key(xa)\n",
        "        v = attention_module.value(xa)\n",
        "        kv_cache_new = (None, None)\n",
        "\n",
        "    wv, qk = attention_module.qkv_attention(q, k, v, mask)\n",
        "    return attention_module.out(wv), kv_cache_new\n",
        "\n",
        "\n",
        "def block_forward(\n",
        "    residual_block,\n",
        "    x: torch.Tensor,\n",
        "    xa: Optional[torch.Tensor] = None,\n",
        "    mask: Optional[torch.Tensor] = None,\n",
        "    kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,\n",
        "):\n",
        "    \"\"\"\n",
        "    Override for residual block forward method for providing kv_cache to attention module.\n",
        "      Parameters:\n",
        "        residual_block: current residual block.\n",
        "        x: input token_ids.\n",
        "        xa: input audio features (Optional).\n",
        "        mask: attention mask (Optional).\n",
        "        kv_cache: cache for storing attention key values.\n",
        "      Returns:\n",
        "        x: residual block output\n",
        "        kv_cache: updated kv_cache\n",
        "\n",
        "    \"\"\"\n",
        "    x0, kv_cache = residual_block.attn(residual_block.attn_ln(\n",
        "        x), mask=mask, kv_cache=kv_cache)\n",
        "    x = x + x0\n",
        "    if residual_block.cross_attn:\n",
        "        x1, _ = residual_block.cross_attn(\n",
        "            residual_block.cross_attn_ln(x), xa)\n",
        "        x = x + x1\n",
        "    x = x + residual_block.mlp(residual_block.mlp_ln(x))\n",
        "    return x, kv_cache\n",
        "\n",
        "\n",
        "\n",
        "# update forward functions\n",
        "for idx, block in enumerate(model.decoder.blocks):\n",
        "    block.forward = partial(block_forward, block)\n",
        "    block.attn.forward = partial(attention_forward, block.attn)\n",
        "    if block.cross_attn:\n",
        "        block.cross_attn.forward = partial(attention_forward, block.cross_attn)\n",
        "\n",
        "\n",
        "def decoder_forward(decoder, x: torch.Tensor, xa: torch.Tensor, kv_cache: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor]]] = None):\n",
        "    \"\"\"\n",
        "    Override for decoder forward method.\n",
        "    Parameters:\n",
        "      x: torch.LongTensor, shape = (batch_size, <= n_ctx) the text tokens\n",
        "      xa: torch.Tensor, shape = (batch_size, n_mels, n_audio_ctx)\n",
        "           the encoded audio features to be attended on\n",
        "      kv_cache: Dict[str, torch.Tensor], attention modules hidden states cache from previous steps\n",
        "    \"\"\"\n",
        "    if kv_cache is not None:\n",
        "        offset = kv_cache[0][0].shape[1]\n",
        "    else:\n",
        "        offset = 0\n",
        "        kv_cache = [None for _ in range(len(decoder.blocks))]\n",
        "    x = decoder.token_embedding(\n",
        "        x) + decoder.positional_embedding[offset: offset + x.shape[-1]]\n",
        "    x = x.to(xa.dtype)\n",
        "    kv_cache_upd = []\n",
        "\n",
        "    for block, kv_block_cache in zip(decoder.blocks, kv_cache):\n",
        "        x, kv_block_cache_upd = block(x, xa, mask=decoder.mask, kv_cache=kv_block_cache)\n",
        "        kv_cache_upd.append(tuple(kv_block_cache_upd))\n",
        "\n",
        "    x = decoder.ln(x)\n",
        "    logits = (\n",
        "        x @ torch.transpose(decoder.token_embedding.weight.to(x.dtype), 1, 0)).float()\n",
        "\n",
        "    return logits, tuple(kv_cache_upd)\n",
        "\n",
        "\n",
        "\n",
        "# override decoder forward\n",
        "model.decoder.forward = partial(decoder_forward, model.decoder)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-jcuygEqXkb3"
      },
      "outputs": [],
      "source": [
        "tokens = torch.ones((5, 3), dtype=torch.int64)\n",
        "logits, kv_cache = model.decoder(tokens, audio_features, kv_cache=None)\n",
        "\n",
        "tokens = torch.ones((5, 1), dtype=torch.int64)\n",
        "\n",
        "if not WHISPER_DECODER_OV.exists():\n",
        "    decoder_model = ov.convert_model(model.decoder, example_input=(tokens, audio_features, kv_cache))\n",
        "    ov.save_model(decoder_model, WHISPER_DECODER_OV)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SFrawf8fXkb3"
      },
      "source": [
        "The decoder model autoregressively predicts the next token guided by encoder hidden states and previously predicted sequence. This means that the shape of inputs which depends on the previous step (inputs for tokens and attention hidden states from previous step) are dynamic. For efficient utilization of memory, you define an upper bound for dynamic input shapes."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cbzyqfo9Xkb3"
      },
      "source": [
        "### Preparing the Inference Pipeline\n",
        "\n",
        "The image below illustrates the pipeline of video transcribing using the Whisper model.\n",
        "\n",
        "![ch06_diagram01.png](https://raw.githubusercontent.com/PacktPublishing/Learn-OpenAI-Whisper/main/Chapter06/ch06_diagram01.png)\n",
        "\n",
        "To run the PyTorch Whisper model, we just need to call the `model.transcribe(audio, **parameters)` function. We will try to reuse original model pipeline for audio transcribing after replacing the original models with OpenVINO IR versions.\n",
        "\n",
        "In the original PyTorch implementation of Whisper, running the transcription pipeline is as simple as calling the `model.transcribe(audio, **parameters)` function, which handles all the necessary steps internally.\n",
        "\n",
        "To leverage the benefits of the OpenVINO toolkit, we will modify this pipeline by replacing the original PyTorch models with their OpenVINO IR counterparts. By doing so, we can take advantage of the performance optimizations and hardware acceleration capabilities offered by OpenVINO while maintaining the overall structure and functionality of the transcription pipeline.\n",
        "\n",
        "In the following sections, we will dive deeper into each step of the pipeline and demonstrate how to integrate the OpenVINO models seamlessly."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z5TTPkpuXkb4"
      },
      "source": [
        "### Selecting the Inference Device\n",
        "\n",
        "One of the key advantages of the OpenVINO toolkit is its ability to optimize and run inference on a wide range of hardware devices, including CPUs, GPUs, and specialized accelerators. To harness this flexibility, we need to specify the target device on which we want to execute the inference pipeline.\n",
        "\n",
        "In the code cell below, you will find a dropdown menu that allows you to select the desired inference device. The available options are dynamically populated based on the devices supported by your system and the installed OpenVINO runtime.\n",
        "\n",
        "Simply choose the appropriate device from the dropdown list, considering factors such as performance, power consumption, and availability. OpenVINO will then optimize the converted models and execute the inference pipeline on the selected device, ensuring the best possible performance and efficiency.\n",
        "\n",
        "By default, the \"AUTO\" option is selected, which allows OpenVINO to automatically choose the most suitable device based on the available hardware and the model's requirements. However, you can override this behavior by explicitly selecting a specific device from the list.\n",
        "\n",
        "Once you have selected the inference device, the subsequent steps in the pipeline will be executed on that device, taking full advantage of the OpenVINO runtime's optimizations and acceleration capabilities.\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "core = ov.Core()"
      ],
      "metadata": {
        "id": "cMZZc6X_ULhD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YFwJYwqwXkb4"
      },
      "outputs": [],
      "source": [
        "import ipywidgets as widgets\n",
        "\n",
        "device = widgets.Dropdown(\n",
        "    options=core.available_devices + [\"AUTO\"],\n",
        "    value='AUTO',\n",
        "    description='Device:',\n",
        "    disabled=False,\n",
        ")\n",
        "\n",
        "device"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "125WuSPjXkb4"
      },
      "outputs": [],
      "source": [
        "from utils import patch_whisper_for_ov_inference, OpenVINOAudioEncoder, OpenVINOTextDecoder\n",
        "\n",
        "patch_whisper_for_ov_inference(model)\n",
        "\n",
        "model.encoder = OpenVINOAudioEncoder(core, WHISPER_ENCODER_OV, device=device.value)\n",
        "model.decoder = OpenVINOTextDecoder(core, WHISPER_DECODER_OV, device=device.value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RwIk3BbwXkb4"
      },
      "source": [
        "## Running the Video Transcription Pipeline\n",
        "\n",
        "With the Whisper model converted to OpenVINO IR format and the inference device selected, we are now ready to run the video transcription pipeline on our chosen video.\n",
        "\n",
        "For the purpose of this tutorial, we will demonstrate the transcription process using a video from YouTube. In the code cell below, you can enter the URL of the YouTube video you wish to transcribe. Please keep in mind that downloading the video may take some time, depending on the video's length and your internet connection speed.\n",
        "\n",
        "Once the video URL is provided, the code will automatically download the video and save it to the local file system. The downloaded video file will serve as the input for the transcription pipeline.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1rSEhiKLXkb5"
      },
      "outputs": [],
      "source": [
        "import ipywidgets as widgets\n",
        "# VIDEO_LINK = \"https://youtu.be/kgL5LBM-hFI\"\n",
        "VIDEO_LINK = \"https://youtu.be/5bs9XoTac88\"\n",
        "link = widgets.Text(\n",
        "    value=VIDEO_LINK,\n",
        "    placeholder=\"Type link for video\",\n",
        "    description=\"Video:\",\n",
        "    disabled=False\n",
        ")\n",
        "\n",
        "link"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u42YLCJJXkb5"
      },
      "outputs": [],
      "source": [
        "from pytube import YouTube\n",
        "from pathlib import Path\n",
        "\n",
        "print(f\"Downloading video {link.value} started\")\n",
        "\n",
        "output_file = Path(\"downloaded_video.mp4\")\n",
        "yt = YouTube(link.value)\n",
        "yt.streams.get_highest_resolution().download(filename=output_file)\n",
        "print(f\"Video saved to {output_file}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YZTn_WCuXkb5"
      },
      "outputs": [],
      "source": [
        "from utils import get_audio\n",
        "\n",
        "audio, duration = get_audio(output_file)\n",
        "\n",
        "import ipywidgets as widgets\n",
        "widgets.Video.from_file(output_file, loop=False, width=400, height=400)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IhPc92RNXkb5"
      },
      "source": [
        "Select the task for the model:\n",
        "\n",
        "* **transcribe** - generate audio transcription in the source language (automatically detected).\n",
        "* **translate** - generate audio transcription with translation to English language."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from whisper import _MODELS\n",
        "import ipywidgets as widgets\n",
        "\n",
        "model_id = widgets.Dropdown(\n",
        "    options=list(_MODELS),\n",
        "    value='base',\n",
        "    description='Model:',\n",
        "    disabled=False,\n",
        ")\n",
        "\n",
        "model_id"
      ],
      "metadata": {
        "id": "AY7wVs9lqf01"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hxCtWXI4Xkb5"
      },
      "outputs": [],
      "source": [
        "task = widgets.Select(\n",
        "    options=[\"transcribe\", \"translate\"],\n",
        "    value=\"translate\",\n",
        "    description=\"Select task:\",\n",
        "    disabled=False\n",
        ")\n",
        "task"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.is_available()"
      ],
      "metadata": {
        "id": "Ei3b5QupqNAU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qC5EHH8MXkb5"
      },
      "outputs": [],
      "source": [
        "transcription = model.transcribe(audio, fp16=torch.cuda.is_available(), task=task.value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6O3UvtmXkb5"
      },
      "source": [
        "\"The results will be saved in the `downloaded_video.srt` file. SRT is one of the most popular formats for storing subtitles and is compatible with many modern video players. This file can be used to embed transcription into videos during playback or by injecting them directly into video files using `ffmpeg`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6-wtpXA6Xkb6"
      },
      "outputs": [],
      "source": [
        "from utils import prepare_srt\n",
        "\n",
        "srt_lines = prepare_srt(transcription, filter_duration=duration)\n",
        "# save transcription\n",
        "with output_file.with_suffix(\".srt\").open(\"w\") as f:\n",
        "    f.writelines(srt_lines)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Sh2y9o1Xkb6"
      },
      "source": [
        "Now let us see the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": [],
        "id": "gJ1ceKJSXkb6"
      },
      "outputs": [],
      "source": [
        "print(\"\".join(srt_lines))"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# prompt: command that creates directy path '/mnt/gradio' if it does not exists\n",
        "\n",
        "!mkdir -p /tmp/gradio"
      ],
      "metadata": {
        "id": "S3lhBkEPmyOL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import gradio as gr\n",
        "\n",
        "# Define a function that Gradio will use to process inputs\n",
        "def video_with_srt(t_video, t_srt):\n",
        "    # Since Gradio handles file paths for video and SRT directly, simply return them\n",
        "    # Note: This assumes t_video and t_srt are paths to the uploaded files\n",
        "    return t_video, t_srt\n",
        "\n",
        "# Create the Gradio interface\n",
        "demo = gr.Interface(\n",
        "    fn=video_with_srt,  # Pass the function reference\n",
        "    inputs=[\n",
        "        gr.Textbox(label=\"Video File Path\"),\n",
        "        gr.Textbox(label=\"SRT File Path\")\n",
        "    ],\n",
        "    outputs=\"video\",  # Specify the output type; this might need adjustment based on the actual handling\n",
        "    examples=[['downloaded_video.mp4', 'downloaded_video.srt']],  # Example inputs\n",
        "    allow_flagging=\"never\"\n",
        ")\n",
        "\n",
        "try:\n",
        "    demo.launch(debug=True)\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "    demo.launch(share=True, debug=True)\n"
      ],
      "metadata": {
        "id": "dxd5UM-vB-ZW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NnZatuyMXkb6"
      },
      "source": [
        "## Interactive Demo\n",
        "\n",
        "To showcase the power and flexibility of the OpenVINO-optimized Whisper model for video transcription, we have created an interactive demo using the Gradio library. This demo allows you to input a YouTube video URL and select the desired transcription task (transcribe or translate) directly from the user interface.\n",
        "\n",
        "Behind the scenes, the demo application downloads the specified video, extracts the audio, and feeds it into the Whisper model for processing. The generated transcription is then displayed in real-time, providing a seamless and user-friendly experience.\n",
        "\n",
        "The code for the interactive demo is provided in the following cells. It includes the necessary setup, such as creating the Gradio interface, defining the transcription function, and configuring the input and output components.\n",
        "\n",
        "Feel free to explore the demo and experiment with different videos and transcription tasks. This interactive component aims to demonstrate the practical application of the OpenVINO-optimized Whisper model and showcase its potential for real-world use cases.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# prompt: command that creates directy path '/mnt/gradio' if it does not exists\n",
        "\n",
        "!mkdir -p /tmp/gradio\n"
      ],
      "metadata": {
        "id": "b6E-tBskZ67T"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "test_replace": {
          "    demo.launch(debug=True)": "    demo.launch()",
          "    demo.launch(share=True, debug=True)": "    demo.launch(share=True)"
        },
        "id": "x3WixdbbXkb7"
      },
      "outputs": [],
      "source": [
        "import gradio as gr\n",
        "from pytube import YouTube\n",
        "from utils import prepare_srt\n",
        "from utils import get_audio\n",
        "\n",
        "def transcribe(url, task):\n",
        "    output_file = Path(\"downloaded_video.mp4\")\n",
        "    yt = YouTube(url)\n",
        "    yt.streams.get_highest_resolution().download(filename=output_file)\n",
        "    audio, duration = get_audio(output_file)\n",
        "    transcription = model.transcribe(audio, fp16=torch.cuda.is_available(), task=task.lower())\n",
        "    srt_lines = prepare_srt(transcription, duration)\n",
        "    with output_file.with_suffix(\".srt\").open(\"w\") as f:\n",
        "        f.writelines(srt_lines)\n",
        "    return [str(output_file), str(output_file.with_suffix(\".srt\"))]\n",
        "\n",
        "\n",
        "demo = gr.Interface(\n",
        "    transcribe,\n",
        "    [gr.Textbox(label=\"YouTube URL\"), gr.Radio([\"Transcribe\", \"Translate\"], value=\"Transcribe\")],\n",
        "    \"video\",\n",
        "    examples=[[\"https://youtu.be/5bs9XoTac88\", \"Translate\"],\n",
        "              [\"https://youtu.be/kgL5LBM-hFI\", \"Transcribe\"]],\n",
        "    allow_flagging=\"never\"\n",
        ")\n",
        "try:\n",
        "    demo.launch(debug=True)\n",
        "except Exception:\n",
        "    demo.launch(share=True, debug=True)\n",
        "# if you are launching remotely, specify server_name and server_port\n",
        "# demo.launch(server_name='your server name', server_port='server port in int')\n",
        "# Read more in the docs: https://gradio.app/docs/"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    },
    "openvino_notebooks": {
      "imageUrl": "https://user-images.githubusercontent.com/29454499/204548693-1304ef33-c790-490d-8a8b-d5766acb6254.png",
      "tags": {
        "categories": [
          "Model Demos",
          "AI Trends"
        ],
        "libraries": [],
        "other": [],
        "tasks": [
          "Speech Recognition"
        ]
      }
    },
    "vscode": {
      "interpreter": {
        "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
      }
    },
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "machine_shape": "hm"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}